# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time

import numpy as np
import torch

from constants import ASR_SAMPLE_RATE
from utils import logger


class VADModel:
    """Voice activity detection."""
    def __init__(self, cfg):
        self.model, _ = torch.hub.load(repo_or_dir=cfg.get('vad_model_path'),
                                       source='local',
                                       model='silero_vad',
                                       force_reload=True,
                                       onnx=True)
        self.warmup()

    def warmup(self):
        for _ in range(10):
            self.model(torch.randn(1, ASR_SAMPLE_RATE), ASR_SAMPLE_RATE)
        self.reset()
        logger.debug('VAD warmup done!')

    def reset(self):
        self.model.reset_states()

    def __call__(self, x):
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x)
        if not isinstance(x, torch.Tensor):
            raise ValueError('Invalid input for VAD module. '
                             'The input requires a numpy array or torch tensor.')

        start = time.time()
        prob = self.model(x, ASR_SAMPLE_RATE).item()
        logger.debug(f'VAD infer time: {(time.time() - start) * 1000:.2f} ms')
        logger.debug(f'VAD prob: {prob:.2f}')
        return prob
